package com.demo.config.sharding;

import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;

import java.util.Collection;
import java.util.HashSet;
import java.util.Set;

/**
 * 范围分片策略
 */
public class RangeShardingStrategy implements RangeShardingAlgorithm<Long> {

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<Long> shardingValue) {

        Range<Long> range = shardingValue.getValueRange();
        Long lowerEndpoint = range.hasLowerBound() ? range.lowerEndpoint() : 0L;
        Long upperEndpoint = range.hasUpperBound() ? range.upperEndpoint() : Long.MAX_VALUE;

        Set<String> database = new HashSet<>();
        // 根据范围值进行遍历，计算需要查询的db
        for (long tmp = lowerEndpoint; tmp <= upperEndpoint; tmp ++) {
            long mod = Math.abs(tmp % ShardingConstant.DB_NUM);
            String tmpDbName = ShardingConstant.DB_NAME + mod;

            // 集合里的db和availableTargetNames长度相同，证明已经把所有db都遍历完
            // 所以结束遍历
            if (database.size() == availableTargetNames.size()) {
                break;
            }

            if (database.contains(tmpDbName)) {
                continue;
            }

            if (availableTargetNames.contains(tmpDbName)) {
                database.add(tmpDbName);
            }
        }
        return database;
    }
}
